%% Manylabs hetero nu-i analysis
% Replicates Figure 7, Table 2

%% Parametrization
% Number of bins for PIT histograms
bin_num = 5;

% Set max iteration number and tolerance
max_iter = 5000;
tol = 0.0001;

% Random seed
seed = 101;

% Sample size of simulated coverage & S stats
sim_cov = 5000;
sim_S = 5000;

% Number of MCMC draws
num_mcmc = 10000;

%% Case 1: N=66, J=5 (All effects included)
N = 66;
J = 5;
highlight_indices = [37:42, 49:60];
EB_est_nu_new_1 = 0.0656581444545457;

% Extract study estimates and ses
sheet1 = 'N=66, J=5';
estimates1 = xlsread('../Data/ManyLabs_Data.xlsx', sheet1, 'C3:H68'); 
se1 = xlsread('../Data/ManyLabs_Data.xlsx', sheet1, 'N3:S68');   

% Randomly assign baseline and validation studies
[hat_theta_1, hat_vartheta_1, rep_base_se2_1, rep_val_se2_1] = random_assignment(estimates1, se1, seed);

% Set prior for nu
theta_var_1 = zeros(N, 1);
for i = 1:N
    theta_var_1(i) = var(hat_theta_1(i,:), 0, 2); 
end
avg_theta_var_1 = mean(theta_var_1);
prior_mean_1 = avg_theta_var_1;

% Generate predictions for validation study estimates
[pred_intervals_1, empirical_cov_1, percentiles_1, nu_draws_1, accept_rate_1] = rwmh_posterior_sampler(hat_theta_1, rep_base_se2_1, hat_vartheta_1, rep_val_se2_1, 0.2, num_mcmc, prior_mean_1, seed);

% Plot 90% CI for nu-i
xtick = 0:10:70;
ytick = 0:0.2:1.4;
[nu_plot_1, nu_plot_order] = plot_nu_CI(nu_draws_1, EB_est_nu_new_1, highlight_indices, xtick, ytick);
saveas(nu_plot_1, "../Results/het_nu_CI_N_66_J_5.png")

% Calculate coverage p-values
cov_stat_sim = simulate_cov_stat_distribution(N, sim_cov);
cov_stat_obs_1 = (empirical_cov_1 - 0.8)^2;
cov_p_value_1 = mean(cov_stat_sim >= cov_stat_obs_1);

%% Case 3: N=48, J=5 (three irreplicable studies excluded)
N = 48;
J = 5;
highlight_indices = [];
EB_est_nu_new_3 = 0.0989465075726676;

% Extract study estimates and ses
sheet3 = 'N=48, J=5';
estimates3 = xlsread('../Data/ManyLabs_Data.xlsx', sheet3, 'C3:H50'); 
se3 = xlsread('../Data/ManyLabs_Data.xlsx', sheet3, 'N3:S50');   

% Randomly assign baseline and validation studies
[hat_theta_3, hat_vartheta_3, rep_base_se2_3, rep_val_se2_3] = random_assignment(estimates3, se3, seed);

% Set prior for nu
theta_var_3 = zeros(N, 1);
for i = 1:N
    theta_var_3(i) = var(hat_theta_3(i,:), 0, 2);
end
avg_theta_var_3 = mean(theta_var_3);
prior_mean_3 = avg_theta_var_3;

% Generate predictions for validation study estimates
[pred_intervals_3, empirical_cov_3, percentiles_3, nu_draws_3, accept_rate_3] = rwmh_posterior_sampler(hat_theta_3, rep_base_se2_3, hat_vartheta_3, rep_val_se2_3, 0.2, num_mcmc, prior_mean_3, seed);

% Plot 90% CI for nu-i
xtick = 0:10:50;
ytick = 0:0.2:1.4;
[nu_plot_3, nu_plot_order] = plot_nu_CI(nu_draws_3, EB_est_nu_new_3, highlight_indices, xtick, ytick);
saveas(nu_plot_3, "../Results/het_nu_CI_N_48_J_5.png")

% Calculate coverage p-values
cov_stat_sim = simulate_cov_stat_distribution(N, sim_cov);
cov_stat_obs_3 = (empirical_cov_3 - 0.8)^2;
cov_p_value_3 = mean(cov_stat_sim >= cov_stat_obs_3);

%% Summarize the EB results
EB_empirical_cov = [empirical_cov_1, empirical_cov_3]
EB_cov_p_value = [cov_p_value_1, cov_p_value_3]
